import datetime
from typing import List, Dict, Any, Optional, Tuple, Set
import numpy as np
import pandas as pd  # ExperimentData 使用 pandas
import networkx as nx  # 用于构建 MemoryGraph


from source.data_processor.utils.experiment_data import ExperimentData
from source.localize.graph.casualGraph import CausalGraph  # 注意: 通常类名是 CausalGraph
from .model.graph import Node, MemoryGraph
from .model.case import Case, CaseData  # SimCase.data 会创建它, 明确导入 Case
from .experiment.simulation import SimCase


class ExperimentCaseConverter:
    """
    将 ExperimentData 和 CausalGraph 转换为 Circa 使用的 Case 对象。
    """

    def __init__(
        self,
        experiment_data: ExperimentData,
        causal_graph: CausalGraph,  # 通常类名是 CausalGraph
        ground_truth_causes: Set[Tuple[str, str]],
    ):
        self.experiment_data = (
            experiment_data  # 实验数据，包含指标时间序列 (metrics_df) 和异常时间戳
        )
        self.causal_graph = causal_graph  # 专家知识或先前生成的因果图
        self.ground_truth_causes = ground_truth_causes  # 真实的根本原因，用于评估等
        self._all_nodes_map: Optional[Dict[Node, int]] = (
            None  # 缓存：Node对象到其在数组中列索引的映射
        )
        self._ordered_nodes: Optional[List[Node]] = (
            None  # 缓存：排序后的Node对象列表，定义了数据数组的列顺序
        )

    def _prepare_time_series_data(
        self,
    ) -> Tuple[np.ndarray, np.ndarray, datetime.timedelta, int]:
        """
        从 ExperimentData.metrics_df 准备时间序列数据。
        时间统一处理为上海时间。

        返回:
            - data_array (np.ndarray): 标准化后的数据（在此代码中是 'value' 列），形状为 (num_timesteps, num_nodes)
            - original_data_array (np.ndarray): 原始未标准化的数据（在此代码中是 'original_value' 列），形状为 (num_timesteps, num_nodes)
            - interval (datetime.timedelta): 采样间隔
            - length_normal (int): 正常期的数据点数量 (到故障时间点为止)
        """
        # 0. 如果 metrics_df 为空，提前处理
        if self.experiment_data.metrics_df.empty:
            # 即使数据为空，也需要获取节点映射，以确定 num_nodes
            ordered_nodes, _ = self._get_or_create_node_mapping()
            num_nodes = len(ordered_nodes)
            # 返回一个符合形状的空数据或默认值
            # 注意：如果 num_nodes 为0，np.zeros((1,0)) 是有效的，但下游可能期望至少一个节点
            return (
                np.zeros((1, num_nodes if num_nodes > 0 else 1)),
                np.zeros((1, num_nodes if num_nodes > 0 else 1)),
                datetime.timedelta(minutes=1),
                0,
            )

        df = self.experiment_data.metrics_df.copy()  # 创建副本以避免修改原始数据

        # 1. 统一时间处理为上海时间
        # 转换 metrics_df 中的 'timestamp' 列
        df["timestamp"] = pd.to_datetime(df["timestamp"])  # 确保是 datetime 对象
        if df["timestamp"].dt.tz is None:
            # 如果是 naive datetime (无时区信息)，假定是 UTC，然后转换为上海时间
            df["timestamp"] = (
                df["timestamp"].dt.tz_localize("UTC").dt.tz_convert("Asia/Shanghai")
            )
        else:
            # 如果已经有时区信息，直接转换为上海时间
            df["timestamp"] = df["timestamp"].dt.tz_convert("Asia/Shanghai")

        # 将 anomaly_timestamp (通常是 Unix UTC 秒级时间戳) 转换为上海时区的 datetime 对象
        anomaly_datetime_sh = pd.to_datetime(
            self.experiment_data.anomaly_timestamp,
            unit="s",
            utc=True,  # 从Unix时间戳（UTC）创建
        ).tz_convert(
            "Asia/Shanghai"
        )  # 转换为上海时间

        # 2. 计算采样间隔 (基于上海时间)
        # 先按 timestamp 排序并去重，用于计算 interval
        # .drop_duplicates() 之后 .sort_values() 更保险，或者先排序再 drop_duplicates()
        sorted_unique_timestamps = df["timestamp"].drop_duplicates().sort_values()
        if len(sorted_unique_timestamps) > 1:
            # 使用 diff() 计算相邻时间戳的差值，然后取中位数作为间隔，这样更稳健，能抵抗少数异常间隔
            interval_seconds = sorted_unique_timestamps.diff().median().total_seconds()
            if pd.isna(interval_seconds) or interval_seconds <= 0:
                # 如果计算出的间隔无效（例如只有一个独特时间戳，diff结果为NaN），则回退到默认值
                interval_seconds = 60  # 默认60秒
        else:
            interval_seconds = 60  # 如果只有一个或没有时间戳，则使用默认间隔
        interval = datetime.timedelta(seconds=interval_seconds)

        # 3. 获取节点映射 (Node对象列表和Node到索引的字典)
        # 这会确定最终输出 NumPy 数组的列顺序和数量
        ordered_nodes, node_to_idx = self._get_or_create_node_mapping()
        num_nodes = len(ordered_nodes)

        # 如果在获取节点映射后发现没有节点（例如，metrics_df 和因果图都为空），特殊处理
        if num_nodes == 0:
            return (
                np.zeros((1, 1)),
                np.zeros((1, 1)),
                interval,
                0,
            )  # 返回(1,1)避免下游维度错误

        # 4. 创建两个 Pivot Table (宽格式数据)，分别用于标准化后的值和原始值
        # 'timestamp' 作为行索引, ('cmdb_id', 'metric_name') 作为多级列索引
        # 'value' 列通常是经过某种标准化/归一化处理的值
        pivot_df_standardized = df.pivot_table(
            index="timestamp",  # 行索引，已经是上海时区
            columns=["cmdb_id", "metric_name"],  # 列索引
            values="value",  # 填充表格的值
        )
        # 'original_value' 列通常是原始的、未经处理的度量值
        pivot_df_original = df.pivot_table(
            index="timestamp",  # 行索引，已经是上海时区
            columns=["cmdb_id", "metric_name"],  # 列索引
            values="original_value",  # 填充表格的值
        )

        # 创建一个完整的多级列索引，基于 _get_or_create_node_mapping() 得到的所有节点
        # 这确保了即使某些节点在当前 metrics_df 中没有数据，它们也会作为列出现在最终数组中
        multi_index_cols = pd.MultiIndex.from_tuples(
            [
                (node.entity, node.metric) for node in ordered_nodes
            ],  # 从Node对象创建元组
            names=["cmdb_id", "metric_name"],  # 列索引的名称
        )

        # Reindex 以确保所有节点列都存在，并按 ordered_nodes 顺序排列
        # 如果 pivot_table 的结果中缺少 multi_index_cols 中的某些列（即某些节点在metrics_df中无数据），
        # reindex 会添加这些列，并用 NaN 填充。然后 .fillna(0) 将这些 NaN 替换为 0。
        reindexed_pivot_df_standardized = pivot_df_standardized.reindex(
            columns=multi_index_cols
        ).fillna(0)
        reindexed_pivot_df_original = pivot_df_original.reindex(
            columns=multi_index_cols
        ).fillna(0)

        data_array = reindexed_pivot_df_standardized.values  # 转换为 NumPy 数组
        original_data_array = reindexed_pivot_df_original.values  # 转换为 NumPy 数组

        # 5. 计算 length_normal (正常期数据点的数量，基于上海时间)
        # reindexed_pivot_df_standardized.index 是排序后的上海时区 DatetimeIndex
        pivot_timestamps_sh = reindexed_pivot_df_standardized.index
        # 计算严格在故障注入时间点 (anomaly_datetime_sh) 之前的数据点数量
        length_normal = (pivot_timestamps_sh < anomaly_datetime_sh).sum()

        # 6. Reshape data_array 和 original_data_array (确保二维)
        # .values 属性通常返回一个至少是二维的数组，但如果原始DataFrame只有一行或一列，
        # 或者num_nodes=1或时间点=1，需要小心处理。
        # 以下的 reshape 逻辑是为了处理极端情况，例如只有一个时间点或只有一个节点。
        # 当前 pivot_table 和 reindex 的方式应该能保证正确的二维形状，
        # 但保留这个检查可以增加代码的稳健性。

        # 如果只有一个时间戳，但有多个节点
        if data_array.ndim == 1 and num_nodes > 1 and len(pivot_timestamps_sh) == 1:
            data_array = data_array.reshape(1, num_nodes)
            original_data_array = original_data_array.reshape(1, num_nodes)
        # 如果有多个时间戳，但只有一个节点
        elif data_array.ndim == 1 and num_nodes == 1 and len(pivot_timestamps_sh) > 1:
            data_array = data_array.reshape(len(pivot_timestamps_sh), 1)
            original_data_array = original_data_array.reshape(
                len(pivot_timestamps_sh), 1
            )
        # 如果只有一个时间戳且只有一个节点
        elif data_array.ndim == 1 and num_nodes == 1 and len(pivot_timestamps_sh) == 1:
            data_array = data_array.reshape(1, 1)
            original_data_array = original_data_array.reshape(1, 1)
        # 如果 data_array.ndim已经是2，则不需要reshape

        return data_array, original_data_array, interval, length_normal

    def _get_or_create_node_mapping(self) -> Tuple[List[Node], Dict[Node, int]]:
        """
        从 metrics_df 中提取所有唯一的 Node 对象，并创建 Node 到列索引的映射。
        同时也确保 CausalGraph 中的节点被包含。
        结果会被缓存。
        """
        # 检查缓存
        if self._ordered_nodes is not None and self._all_nodes_map is not None:
            return self._ordered_nodes, self._all_nodes_map

        unique_nodes: Set[Node] = set()  # 使用集合来存储唯一的Node对象

        # 1. 从 metrics_df 中提取节点
        if not self.experiment_data.metrics_df.empty:
            # 优化：避免 iterrows()，直接操作列，然后去重
            temp_df = self.experiment_data.metrics_df[
                ["cmdb_id", "metric_name"]
            ].drop_duplicates()
            for (
                _,
                row,
            ) in temp_df.iterrows():  # 在去重后的少量数据上 iterrows() 性能尚可
                unique_nodes.add(Node(entity=row["cmdb_id"], metric=row["metric_name"]))

        # 2. 从 CausalGraph 的节点定义中提取节点
        if (
            self.causal_graph and self.causal_graph.nodes
        ):  # 检查 causal_graph及其nodes属性是否存在
            for cg_node_id, cg_node_obj in self.causal_graph.nodes.items():
                entity, metric = None, None
                # 尝试从节点对象属性中获取 entity_id 和 metric_name
                if hasattr(cg_node_obj, "entity_id") and hasattr(
                    cg_node_obj, "metric_name"
                ):
                    entity = cg_node_obj.entity_id
                    metric = cg_node_obj.metric_name
                # 如果属性不存在，尝试从节点ID（通常是字符串）中解析
                # 假设节点ID格式为 "entity|metric"
                elif isinstance(cg_node_id, str):
                    try:
                        parts = cg_node_id.split("|", 1)
                        if len(parts) == 2:
                            entity, metric = parts[0], parts[1]
                    except ValueError:
                        # 解析失败，可以选择跳过或记录日志
                        pass  # print(f"Warning: Could not parse node_id '{cg_node_id}' from causal_graph.nodes")
                if entity and metric:
                    unique_nodes.add(Node(entity=entity, metric=metric))

        # 3. 从 CausalGraph 的边定义中提取节点
        if (
            self.causal_graph and self.causal_graph.causal_edges
        ):  # 检查 causal_graph及其causal_edges属性是否存在
            for edge in self.causal_graph.causal_edges:
                # 确保边对象有预期的属性
                if hasattr(edge, "cause_pod") and hasattr(edge, "cause_metric"):
                    unique_nodes.add(
                        Node(entity=edge.cause_pod, metric=edge.cause_metric)
                    )
                if hasattr(edge, "effect_pod") and hasattr(edge, "effect_metric"):
                    unique_nodes.add(
                        Node(entity=edge.effect_pod, metric=edge.effect_metric)
                    )

        # 对所有收集到的唯一节点进行排序，以确保每次生成的列顺序一致
        self._ordered_nodes = sorted(
            list(unique_nodes), key=lambda n: (n.entity, n.metric)
        )
        # 创建 Node 对象到其在排序后列表中的索引的映射
        self._all_nodes_map = {node: i for i, node in enumerate(self._ordered_nodes)}

        return self._ordered_nodes, self._all_nodes_map

    def _build_memory_graph(self) -> MemoryGraph:
        """
        根据 CausalGraph 构建 MemoryGraph。
        MemoryGraph 使用 networkx.DiGraph 作为底层图结构。
        """
        ordered_nodes, _ = self._get_or_create_node_mapping()  # 获取所有节点的标准顺序
        nx_graph = nx.DiGraph()  # 创建一个有向图实例

        # 首先将所有在时间序列数据或因果图定义中出现过的节点添加到图中
        # 这样可以确保即使某些节点是孤立的（没有边），它们也存在于图中
        nx_graph.add_nodes_from(ordered_nodes)  # add_nodes_from 可以接受一个节点列表

        # 如果存在因果图的边定义，则将它们添加到 nx_graph 中
        if self.causal_graph and self.causal_graph.causal_edges:
            for c_edge in self.causal_graph.causal_edges:
                # 从因果边对象中提取源节点和目标节点的entity和metric
                cause_node = Node(entity=c_edge.cause_pod, metric=c_edge.cause_metric)
                effect_node = Node(
                    entity=c_edge.effect_pod, metric=c_edge.effect_metric
                )

                # 获取边的属性，如 p_value, strength, time_lag
                # 使用 getattr 提供默认值 None，以防属性不存在
                p_value = getattr(c_edge, "p_value", None)
                strength = getattr(c_edge, "strength", None)
                time_lag = getattr(c_edge, "time_lag", None)

                # 理论上，这些节点应该已经在 nx_graph.add_nodes_from(ordered_nodes) 中被添加
                # 但为了安全起见，可以再次确认或添加（add_node 若节点已存在则无操作）
                # if cause_node not in nx_graph:
                #     nx_graph.add_node(cause_node) # 如果节点可能未被ordered_nodes覆盖，则需要此步
                # if effect_node not in nx_graph:
                #     nx_graph.add_node(effect_node)

                # 添加边，并附带其属性
                nx_graph.add_edge(
                    cause_node,
                    effect_node,
                    p_value=p_value,
                    strength=strength,
                    time_lag=time_lag,
                )
        return MemoryGraph(graph=nx_graph)  # 用构建好的nx图实例化MemoryGraph

    def create_case(
        self,
    ) -> Case:  # 修改返回类型为 Case (SimCase 是 Case 的一个具体实现)
        """
        创建 Case 对象。
        Case 对象封装了所有用于根因定位算法的输入数据。
        """
        # 1. 准备时间序列数据数组、采样间隔和正常期长度
        #   _prepare_time_series_data 内部会调用 _get_or_create_node_mapping，
        #   所以 self._ordered_nodes 会被设置。
        ts_data, original_ts_data, interval, length_normal = (
            self._prepare_time_series_data()
        )

        # 2. 构建记忆图 (MemoryGraph)
        #   MemoryGraph 表示节点间的因果关系知识。
        memory_graph = self._build_memory_graph()

        # 3. 创建 SimCase (或通用 Case) 对象
        #   SimCase 是 Case 的一种，通常用于模拟实验，它包含了真实原因。
        #   确保 self._ordered_nodes 在这里是可用的，_prepare_time_series_data 应该已经设置了它。
        #   如果 _prepare_time_series_data 因为 metrics_df 为空而提前返回，
        #   它内部调用的 _get_or_create_node_mapping 仍然会设置 self._ordered_nodes。
        sim_case = SimCase(
            data=ts_data,  # (标准化)时间序列数据
            original_data=original_ts_data,  # 原始时间序列数据
            causes=self.ground_truth_causes,  # 真实的根本原因集合
            length_normal=length_normal,  # 正常期数据点长度
            interval=interval,  # 采样间隔
            graph=memory_graph,  # 因果关系图 (MemoryGraph)
            nodes=self._ordered_nodes,  # 节点列表，定义了数据数组的列顺序
        )
        return sim_case  # 返回 SimCase 实例
